from rest_framework import viewsets, status
from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework.permissions import IsAuthenticated
from django.db.models import Sum
from .models import ProvincePatent, Keyword, WordPair, TopOrg, WordTrend
from .serializers import (
    ProvincePatentSerializer,
    KeywordSerializer,
    WordPairSerializer,
    TopOrgSerializer
)
import os
from django.conf import settings
from .charts.heatmap import ChinaPatentHeatmap
from datetime import datetime
from .charts.bubble import KeywordBubbleChart
from .charts.barchart import TopInstitutesGenerator
from .charts.linechart import KeywordTrendGenerator
from .charts.scatter import KeywordDistributionGenerator
from users.authentication import TokenAuthentication
from django.utils.text import slugify
import pypinyin
from django.db import models


class TrendViewSet(viewsets.ViewSet):
    """趋势分析视图集"""
    authentication_classes = [TokenAuthentication]
    permission_classes = [IsAuthenticated]

    def get_pinyin_dir(self, industry):
        """获取产业名称的拼音目录名"""
        pinyin_list = pypinyin.lazy_pinyin(industry)
        return ''.join(pinyin_list)

    @action(detail=False, methods=['get'])
    def heatmap(self, request):
        """生成专利分布热力图"""
        try:
            # 获取产业参数
            industry = request.query_params.get('industry')
            if not industry:
                return Response({
                    'code': 400,
                    'message': '请指定产业名称'
                }, status=400)

            # 获取指定产业的专利数据
            patents = ProvincePatent.objects.filter(industry=industry)
            if not patents.exists():
                return Response({
                    'code': 200,
                    'dataExists': False,
                    'message': f'没有找到{industry}的专利数据'
                })

            # 生成产业拼音目录
            pinyin_dir = self.get_pinyin_dir(industry)

            # 使用绝对路径
            output_dir = os.path.join(settings.BASE_DIR, 'trend', 'output', pinyin_dir)
            os.makedirs(output_dir, exist_ok=True)

            # 生成热力图
            heatmap = ChinaPatentHeatmap()
            filename = 'heatmap.html'
            save_path = os.path.join(output_dir, filename)

            # 生成图表
            heatmap.generate(patents, save_path)

            # 检查文件是否生成
            if not os.path.exists(save_path):
                print(f"File not found: {save_path}")  # 添加调试信息
                return Response({
                    'code': 500,
                    'message': '热力图生成失败'
                }, status=500)

            # 使用相对URL路径
            static_url = f'/static/trend/output/{pinyin_dir}/{filename}'

            print(f"File saved at: {save_path}")  # 添加调试信息
            print(f"Static URL: {static_url}")  # 添加调试信息

            return Response({
                'code': 200,
                'dataExists': True,
                'data': {
                    'imageUrl': static_url,
                    'distribution': {
                        item.province: item.count
                        for item in patents
                    }
                },
                'message': 'success'
            })

        except Exception as e:
            print(f"Error generating heatmap: {str(e)}")  # 添加错误日志
            return Response({
                'code': 500,
                'message': str(e)
            }, status=500)

    @action(detail=False, methods=['get'])
    def bubble(self, request):
        """生成核心词共现关系气泡图"""
        try:
            # 获取产业参数
            industry = request.query_params.get('industry')
            if not industry:
                return Response({
                    'code': 400,
                    'message': '请指定产业名称'
                }, status=400)

            # 获取关键词数据（仅取前20个关键词）
            keywords = dict(
                Keyword.objects.filter(industry=industry)
                .order_by('-count')  # 按计数降序排序
                .values_list('keyword', 'count')[:20]  # 只取前20个
            )

            # 添加调试信息
            print(f"Keywords data: {keywords}")

            # 获取词对关系数据
            word_pairs = WordPair.objects.select_related('kword1', 'kword2').filter(
                kword1__industry=industry,
                kword1__keyword__in=keywords.keys(),  # 只取与前20个关键词相关的词对
                kword2__keyword__in=keywords.keys()  # 确保词对中的关键词都在前20个中
            ).values_list('kword1__keyword', 'kword2__keyword', 'co_count')

            # 添加调试信息
            print(f"Word pairs data: {list(word_pairs)}")

            if not keywords:
                return Response({
                    'code': 200,
                    'dataExists': False,
                    'message': f'没有找到{industry}的关键词数据'
                })

            if not word_pairs:
                return Response({
                    'code': 200,
                    'dataExists': False,
                    'message': f'没有找到{industry}的词对关系数据'
                })

            # 生成产业拼音目录
            pinyin_dir = self.get_pinyin_dir(industry)
            output_dir = os.path.join(settings.BASE_DIR, 'trend', 'output', pinyin_dir)
            os.makedirs(output_dir, exist_ok=True)

            # 生成气泡图
            bubble = KeywordBubbleChart()
            filename = 'bubble.html'
            save_path = os.path.join(output_dir, filename)

            bubble.generate(keywords, word_pairs, save_path)

            if not os.path.exists(save_path):
                print(f"File not found: {save_path}")
                return Response({
                    'code': 500,
                    'message': '气泡图生成失败'
                }, status=500)

            static_url = f'/static/trend/output/{pinyin_dir}/{filename}'
            print(f"File saved at: {save_path}")
            print(f"Static URL: {static_url}")

            return Response({
                'code': 200,
                'dataExists': True,
                'data': {
                    'imageUrl': static_url,
                    'keywords': [
                        {'word': word, 'count': count}
                        for word, count in keywords.items()
                    ],
                    'relations': [
                        {
                            'source': w1,
                            'target': w2,
                            'count': count
                        }
                        for w1, w2, count in word_pairs
                    ]
                },
                'message': 'success'
            })

        except Exception as e:
            print(f"Error generating bubble chart: {str(e)}")
            return Response({
                'code': 500,
                'message': str(e)
            }, status=500)

    @action(detail=False, methods=['get'])
    def barchart(self, request):
        """生成机构专利排名柱状图"""
        try:
            # 获取产业参数
            industry = request.query_params.get('industry')
            if not industry:
                return Response({
                    'code': 400,
                    'message': '请指定产业名称'
                }, status=400)

            # 获取指定产业的机构数据
            orgs = TopOrg.objects.filter(industry=industry).order_by('-count')[:10]
            if not orgs.exists():
                return Response({
                    'code': 200,
                    'dataExists': False,
                    'message': f'没有找到{industry}的机构数据'
                })

            # 生成产业拼音目录
            pinyin_dir = self.get_pinyin_dir(industry)
            output_dir = os.path.join(os.path.dirname(__file__), 'output', pinyin_dir)
            os.makedirs(output_dir, exist_ok=True)

            # 转换数据格式
            data = {org.org: org.count for org in orgs}

            # 生成柱状图
            generator = TopInstitutesGenerator()
            save_path = os.path.join(output_dir, 'barchart.html')
            generator.generate(data, save_path)

            return Response({
                'code': 200,
                'dataExists': True,
                'data': {
                    'imageUrl': f'/static/trend/output/{pinyin_dir}/barchart.html',
                    'institutes': [
                        {'name': org.org, 'count': org.count}
                        for org in orgs
                    ]
                },
                'message': 'success'
            })

        except Exception as e:
            return Response({
                'code': 500,
                'message': str(e)
            }, status=500)

    @action(detail=False, methods=['get'])
    def linechart(self, request):
        """生成核心词趋势折线图"""
        try:
            # 获取产业参数
            industry = request.query_params.get('industry')
            if not industry:
                return Response({
                    'code': 400,
                    'message': '请指定产业名称'
                }, status=400)

            # 1. 先获取2023年的数据并排序
            top_words_2023 = WordTrend.objects.filter(
                industry=industry,
                year=2023
            ).order_by('-patent_count')[:5].values_list('word', flat=True)

            # 转换为列表
            top_words = list(top_words_2023)

            # 2. 获取这些词的所有年份数据
            trends = WordTrend.objects.filter(
                industry=industry,
                word__in=top_words
            ).order_by('year', 'word')

            if not trends.exists():
                return Response({
                    'code': 200,
                    'dataExists': False,
                    'message': f'没有找到{industry}的趋势数据'
                })

            # 生成产业拼音目录
            pinyin_dir = self.get_pinyin_dir(industry)
            output_dir = os.path.join(os.path.dirname(__file__), 'output', pinyin_dir)
            os.makedirs(output_dir, exist_ok=True)

            # 转换数据格式，并确保每个词都有5年的数据
            data = {}
            years = list(range(2019, 2024))

            # 初始化数据结构
            for word in top_words:
                data[word] = {year: 0 for year in years}  # 默认值为0

            # 填充实际数据
            for trend in trends:
                data[trend.word][trend.year] = trend.patent_count

            # 转换为列表格式
            final_data = {
                word: [counts[year] for year in years]
                for word, counts in data.items()
            }

            # 生成趋势图
            generator = KeywordTrendGenerator()
            save_path = os.path.join(output_dir, 'linechart.html')
            generator.generate(
                data=final_data,
                save_path=save_path,
                title=f'{industry}领域核心词专利趋势(2019-2023)'
            )

            return Response({
                'code': 200,
                'dataExists': True,
                'data': {
                    'imageUrl': f'/static/trend/output/{pinyin_dir}/linechart.html',
                    'trends': {
                        word: {
                            'patent_counts': counts,
                            'years': years
                        }
                        for word, counts in final_data.items()
                    }
                },
                'message': 'success'
            })

        except Exception as e:
            print(f"Error in linechart: {str(e)}")  # 添加错误日志
            return Response({
                'code': 500,
                'message': str(e)
            }, status=500)

    @action(detail=False, methods=['get'])
    def scatter(self, request):
        """生成核心词分布散点图"""
        try:
            # 获取产业参数
            industry = request.query_params.get('industry')
            if not industry:
                return Response({
                    'code': 400,
                    'message': '请指定产业名称'
                }, status=400)

            # 获取前5个关键词的趋势数据
            # 修改查询方式，避免使用子查询
            top_words = WordTrend.objects.filter(
                industry=industry
            ).values('word').annotate(
                total_count=models.Sum('count')
            ).order_by('-total_count')[:5].values_list('word', flat=True)

            # 获取这些词的年度数据
            trends = WordTrend.objects.filter(
                industry=industry,
                word__in=list(top_words)  # 转换为列表
            )

            # 添加调试信息
            print(f"Top words: {list(top_words)}")
            print(f"Trends count: {trends.count()}")

            if not trends.exists():
                return Response({
                    'code': 200,
                    'dataExists': False,
                    'message': f'没有找到{industry}的趋势数据'
                })

            # 转换数据格式
            trend_data = {}
            for trend in trends:
                if trend.word not in trend_data:
                    trend_data[trend.word] = {}
                trend_data[trend.word][trend.year] = trend.patent_count
            print(f"Trend data: {trend_data}")
            # 生成产业拼音目录
            pinyin_dir = self.get_pinyin_dir(industry)
            output_dir = os.path.join(settings.BASE_DIR, 'trend', 'output', pinyin_dir)
            os.makedirs(output_dir, exist_ok=True)

            # 生成散点图
            generator = KeywordDistributionGenerator()
            filename = 'scatter.html'
            save_path = os.path.join(output_dir, filename)

            generator.generate(trend_data, save_path)

            if not os.path.exists(save_path):
                print(f"File not found: {save_path}")
                return Response({
                    'code': 500,
                    'message': '散点图生成失败'
                }, status=500)

            static_url = f'/static/trend/output/{pinyin_dir}/{filename}'
            print(f"File saved at: {save_path}")
            print(f"Static URL: {static_url}")

            return Response({
                'code': 200,
                'dataExists': True,
                'data': {
                    'imageUrl': static_url,
                    'trends': trend_data
                },
                'message': 'success'
            })

        except Exception as e:
            print(f"Error generating scatter chart: {str(e)}")
            return Response({
                'code': 500,
                'message': str(e)
            }, status=500)

    @action(detail=False, methods=['get'])
    def wordcloud(self, request):
        """生成核心词词云图"""
        try:
            # 获取产业参数
            industry = request.query_params.get('industry')
            if not industry:
                return Response({
                    'code': 400,
                    'message': '请指定产业名称'
                }, status=400)

            # 获取年份参数，默认2023
            year = request.query_params.get('year', '2023')
            try:
                year = int(year)
                if year < 2019 or year > 2023:
                    return Response({
                        'code': 400,
                        'message': '年份必须在2019-2023之间'
                    }, status=400)
            except ValueError:
                return Response({
                    'code': 400,
                    'message': '无效的年份格式'
                }, status=400)

            # 获取指定产业和年份的词频数据
            word_trends = WordTrend.objects.filter(
                industry=industry,
                year=year
            ).order_by('-patent_count')[:30]  # 取专利数量前20的词

            if not word_trends.exists():
                return Response({
                    'code': 200,
                    'dataExists': False,
                    'message': f'没有找到{industry}在{year}年的词频数据'
                })

            # 生成产业拼音目录
            pinyin_dir = self.get_pinyin_dir(industry)
            output_dir = os.path.join(os.path.dirname(__file__), 'output', pinyin_dir)
            os.makedirs(output_dir, exist_ok=True)

            # 转换数据格式
            cloud_data = {
                item.word: item.patent_count
                for item in word_trends
            }

            # 生成词云图
            from .charts.wordcloud import KeywordCloudGenerator
            generator = KeywordCloudGenerator(
                min_font_size=12,
                max_font_size=60
            )

            filename = f'wordcloud_{year}.html'
            save_path = os.path.join(output_dir, filename)

            generator.generate(
                data=cloud_data,
                year=year,
                save_path=save_path,
                title=f'{industry}领域{year}年核心词专利分布词云图'
            )

            return Response({
                'code': 200,
                'dataExists': True,
                'data': {
                    'imageUrl': f'/static/trend/output/{pinyin_dir}/{filename}',
                    'year': year,
                    'words': [
                        {
                            'word': item.word,
                            'patent_count': item.patent_count,
                            'count': item.count
                        }
                        for item in word_trends
                    ],
                    'total_words': len(cloud_data),
                    'max_count': max(cloud_data.values()),
                    'min_count': min(cloud_data.values())
                },
                'message': 'success'
            })

        except Exception as e:
            return Response({
                'code': 500,
                'message': str(e)
            }, status=500)